#include "matrix.h"

// 矩阵创建
Matrix *create_matrix(int rows, int cols)
{
    Matrix *mat = (Matrix *)malloc(sizeof(Matrix));
    mat->rows = rows;
    mat->cols = cols;
    mat->data = (double **)malloc(rows * sizeof(double *));
    for (int i = 0; i < rows; i++)
    {
        mat->data[i] = (double *)calloc(cols, sizeof(double));
    }
    printf("mat created\n");
    return mat;
}

// 矩阵释放
void free_matrix(Matrix *mat)
{
    if (mat)
    {
        for (int i = 0; i < mat->rows; i++)
        {
            free(mat->data[i]);
        }
        free(mat->data);
        free(mat);
    }
}

// 矩阵乘法
Matrix *matrix_multiply(Matrix *a, Matrix *b)
{
    if (a->cols != b->rows)
        return NULL;

    Matrix *result = create_matrix(a->rows, b->cols);
    for (int i = 0; i < a->rows; i++)
    {
        for (int j = 0; j < b->cols; j++)
        {
            result->data[i][j] = 0;
            for (int k = 0; k < a->cols; k++)
            {
                result->data[i][j] += a->data[i][k] * b->data[k][j];
            }
        }
    }
    return result;
}

// 矩阵转置
Matrix *matrix_transpose(Matrix *mat)
{
    Matrix *result = create_matrix(mat->cols, mat->rows);
    for (int i = 0; i < mat->rows; i++)
    {
        for (int j = 0; j < mat->cols; j++)
        {
            result->data[j][i] = mat->data[i][j];
        }
    }
    return result;
}

// 矩阵加法
Matrix *matrix_add(Matrix *a, Matrix *b)
{
    if (a->rows != b->rows || a->cols != b->cols)
        return NULL;

    Matrix *result = create_matrix(a->rows, a->cols);
    for (int i = 0; i < a->rows; i++)
    {
        for (int j = 0; j < a->cols; j++)
        {
            result->data[i][j] = a->data[i][j] + b->data[i][j];
        }
    }
    return result;
}

// 矩阵减法
Matrix *matrix_subtract(Matrix *a, Matrix *b)
{
    if (a->rows != b->rows || a->cols != b->cols)
        return NULL;

    Matrix *result = create_matrix(a->rows, a->cols);
    for (int i = 0; i < a->rows; i++)
    {
        for (int j = 0; j < a->cols; j++)
        {
            result->data[i][j] = a->data[i][j] - b->data[i][j];
        }
    }
    return result;
}

// 矩阵数乘
Matrix *matrix_scalar_multiply(Matrix *mat, double scalar)
{
    Matrix *result = create_matrix(mat->rows, mat->cols);
    for (int i = 0; i < mat->rows; i++)
    {
        for (int j = 0; j < mat->cols; j++)
        {
            result->data[i][j] = mat->data[i][j] * scalar;
        }
    }
    return result;
}

// 单位矩阵
Matrix *matrix_eye(int n)
{
    Matrix *result = create_matrix(n, n);
    for (int i = 0; i < n; i++)
    {
        result->data[i][i] = 1.0;
    }
    return result;
}

// 矩阵复制
Matrix *matrix_copy(Matrix *mat)
{
    if (!mat)
        return NULL;

    Matrix *result = create_matrix(mat->rows, mat->cols);
    for (int i = 0; i < mat->rows; i++)
    {
        memcpy(result->data[i], mat->data[i], mat->cols * sizeof(double));
    }
    return result;
}

// 辅助函数：交换矩阵的两行
static void swap_rows(Matrix *mat, int i, int j)
{
    double *temp = mat->data[i];
    mat->data[i] = mat->data[j];
    mat->data[j] = temp;
}

// 辅助函数：交换数组的两个元素
static void swap_elements(double *a, double *b)
{
    double temp = *a;
    *a = *b;
    *b = temp;
}

// 矩阵行列式计算（使用LU分解）
double matrix_det(Matrix *mat)
{
    if (!mat || mat->rows != mat->cols)
        return NAN;

    int n = mat->rows;
    Matrix *temp = matrix_copy(mat);
    double det = 1.0;
    int sign = 1;

    for (int i = 0; i < n; i++)
    {
        // 找到主元
        int pivot = i;
        for (int j = i; j < n; j++)
        {
            if (fabs(temp->data[j][i]) > fabs(temp->data[pivot][i]))
            {
                pivot = j;
            }
        }

        if (pivot != i)
        {
            swap_rows(temp, i, pivot);
            sign *= -1; // 交换行改变行列式符号
        }

        // 如果主元为0，行列式为0
        if (fabs(temp->data[i][i]) < 1e-10)
        {
            free_matrix(temp);
            return 0.0;
        }

        det *= temp->data[i][i];

        // 消去下方元素
        for (int j = i + 1; j < n; j++)
        {
            double factor = temp->data[j][i] / temp->data[i][i];
            for (int k = i; k < n; k++)
            {
                temp->data[j][k] -= factor * temp->data[i][k];
            }
        }
    }

    free_matrix(temp);
    return det * sign;
}

// 矩阵求逆（使用增广矩阵法）
int matrix_inverse(Matrix *mat, Matrix *inv)
{
    if (!mat || !inv || mat->rows != mat->cols ||
        inv->rows != mat->rows || inv->cols != mat->cols)
    {
        return 0; // 失败
    }

    int n = mat->rows;
    Matrix *aug = create_matrix(n, 2 * n); // 增广矩阵 [A|I]

    // 初始化增广矩阵
    for (int i = 0; i < n; i++)
    {
        for (int j = 0; j < n; j++)
        {
            aug->data[i][j] = mat->data[i][j];
        }
        aug->data[i][i + n] = 1.0;
    }

    for (int i = 0; i < n; i++)
    {
        // 找到主元
        int pivot = i;
        for (int j = i; j < n; j++)
        {
            if (fabs(aug->data[j][i]) > fabs(aug->data[pivot][i]))
            {
                pivot = j;
            }
        }

        if (pivot != i)
        {
            swap_rows(aug, i, pivot);
        }

        // 如果主元接近0，矩阵不可逆
        if (fabs(aug->data[i][i]) < 1e-10)
        {
            free_matrix(aug);
            return 0; // 失败
        }

        // 归一化主行
        double div = aug->data[i][i];
        for (int j = i; j < 2 * n; j++)
        {
            aug->data[i][j] /= div;
        }

        // 消去其他行
        for (int j = 0; j < n; j++)
        {
            if (j != i && fabs(aug->data[j][i]) > 1e-10)
            {
                double factor = aug->data[j][i];
                for (int k = i; k < 2 * n; k++)
                {
                    aug->data[j][k] -= factor * aug->data[i][k];
                }
            }
        }
    }

    // 提取逆矩阵
    for (int i = 0; i < n; i++)
    {
        for (int j = 0; j < n; j++)
        {
            inv->data[i][j] = aug->data[i][j + n];
        }
    }

    free_matrix(aug);
    return 1; // 成功
}
